def get_result_dict(file_path):
    out_dict = {}
    infile = open(file_path, 'r')
    lines = infile.readlines()
    for line in lines:
        line_vec = line.strip().split("\t")
        if len(line_vec) != 4:
            continue
        imgname = line_vec[0]
        clas = line_vec[1]
        conf = line_vec[2]
        tmp_bbox = line_vec[3]
        tmp_bbox = tmp_bbox.strip("[").strip("]")
        bbox = tmp_bbox.split(' ')
        # print("tmp_bbox type:",type(tmp_bbox))
        # print("tmp_bbox:",tmp_bbox)
        # bbox = json.loads(tmp_bbox)
        # print("bbox:",bbox)
        float_bbox = [float(item) for item in bbox if len(item) > 0]
        bbox = numpy.array(float_bbox)
        tmpkey = str(clas) + "_" + imgname

        if tmpkey in out_dict:
            out_dict[tmpkey].append([clas, conf, bbox])
        else:
            tmp_list = [[clas, conf, bbox]]
            out_dict[tmpkey] = tmp_list
    print("run end total item:", len(out_dict))
    return out_dict


def get_fround_truth_dict(file_path):
    out_dict = {}
    infile = open(file_path, 'r')
    lines = infile.readlines()
    for line in lines:
        line_vec = line.strip().split("\t")
        if len(line_vec) != 3:
            continue
        imgname = line_vec[0]
        clas = int(line_vec[1])
        tmp_bbox = line_vec[2]
        tmp_bbox = tmp_bbox.strip("[").strip("]")
        bbox = tmp_bbox.split(',')
        # print("tmp_bbox type:",type(tmp_bbox))
        # print("tmp_bbox:",tmp_bbox)
        # bbox = json.loads(tmp_bbox)
        # print("bbox:",bbox)
        # print("bbox:",bbox)
        float_bbox = [float(item) for item in bbox if len(item) > 0]
        bbox = numpy.array(float_bbox)
        # tmpkey = str(clas) + "_" + imgname
        objtype = labels[clas]
        if imgname in out_dict:
            # curdict = out_dict[tmpkey]

            out_dict[imgname][objtype].append([clas, bbox])
            # out_dict[tmpkey].append([clas, bbox])
        else:
            tmp_dict = {labels[0]: [], labels[1]: [], labels[2]: []}
            tmp_dict[objtype].append([clas, bbox])
            out_dict[imgname] = tmp_dict
    print("run end total item:", len(out_dict))
    print("out_dict:", out_dict)
    return out_dict


def generate_detections(yitu_file):
    infile = open(yitu_file)
    lines = infile.readlines()
    print("filename:",yitu_file)
    print("lines:",len(lines))
    count = 0
    for line in lines:
        line_vec = line.strip().split("\t")
        if len(line_vec) != 4:
            continue

        output_file_name = "detections/" + line_vec[0].split("/")[-1].split(".")[0] + ".txt"
        print("filename:", output_file_name)
        clas = line_vec[1]
        conf = line_vec[2]
        bbox = line_vec[3].strip("[").strip("]").split(' ')
        outstr = "{} {} {} {} {} {}\n".format(clas, conf, bbox[0], bbox[1], bbox[2], bbox[3])
        with open(output_file_name, 'a+') as outfile:
            outfile.write(outstr)
            count += 1
    print("run end total:",count)

def generate_groundtruths(ground_truth_file):
    infile = open(ground_truth_file)
    lines = infile.readlines()
    count = 0
    for line in lines:
        line_vec = line.strip().split("\t")
        if len(line_vec) != 3:
            continue

        output_file_name = "groundtruths/" + line_vec[0].split("/")[-1].split(".")[0] + ".txt"
        print("filename:", output_file_name)
        clas = line_vec[1]
        bbox = line_vec[2].strip("[").strip("]").split(' ')
        outstr = "{} {} {} {} {}\n".format(clas, bbox[0], bbox[1], bbox[2], bbox[3])
        with open(output_file_name, 'a+') as outfile:
            outfile.write(outstr)
            count += 1
    print("run end total:",count)

if __name__ == "__main__":
    yitu_file = "result.txt"
    generate_detections(yitu_file)
    ground_truth_file = "nvda_yolo_result.txt"
    generate_groundtruths(ground_truth_file)
